"""
Elasticsearch 客户端封装，支持批量写入和重试机制。
优化版：使用单例模式和异步上下文管理器，避免重复实例化和提高内存管理效率。
"""
from typing import TypeVar, Type, Callable, Optional, List, Dict, Any
from pydantic import BaseModel
from elasticsearch import AsyncElasticsearch
from elasticsearch.helpers import async_bulk
import logging
from typing import Optional, List, Dict, Any
import asyncio
from config import ELASTICSEARCH_HOSTS, ELASTICSEARCH_REQUEST_TIMEOUT, ELASTICSEARCH_USERNAME, ELASTICSEARCH_PASSWORD
from utils.logger_utils import retry

logger = logging.getLogger(__name__)


class ESClient:
    """
    Elasticsearch 客户端单例类，提供异步操作接口
    """
    _instance: Optional["ESClient"] = None
    _lock = asyncio.Lock()
    _client: Optional[AsyncElasticsearch] = None
    _init_count = 0  # 跟踪初始化计数

    def __new__(cls) -> "ESClient":
        if cls._instance is None:
            cls._instance = super(ESClient, cls).__new__(cls)
        return cls._instance

    async def __aenter__(self) -> "ESClient":
        """异步上下文管理器入口"""
        if self._client is None:
            await self.initialize()
        self._init_count += 1
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """异步上下文管理器出口，仅在所有引用结束时关闭客户端"""
        self._init_count -= 1
        if self._init_count <= 0:
            await self.close()

    async def initialize(self) -> None:
        """初始化ES客户端（如果尚未初始化）"""
        async with self._lock:
            if self._client is None:
                logger.info("Initializing Elasticsearch client")
                self._client = AsyncElasticsearch(
                    hosts=ELASTICSEARCH_HOSTS,
                    request_timeout=ELASTICSEARCH_REQUEST_TIMEOUT,
                    http_auth=(ELASTICSEARCH_USERNAME, ELASTICSEARCH_PASSWORD),
                    max_connections=50,
                    retry_on_timeout=True,
                    retry_on_status={502, 503, 504}
                )

    async def close(self) -> None:
        """关闭ES客户端连接"""
        async with self._lock:
            if self._client is not None:
                logger.info("Closing Elasticsearch client")
                await self._client.close()
                self._client = None
                self._init_count = 0

    @property
    def client(self) -> AsyncElasticsearch:
        """获取ES客户端实例，如果未初始化则抛出异常"""
        if self._client is None:
            raise RuntimeError("Elasticsearch client not initialized. Use 'async with ESClient()' context manager.")
        return self._client

    @retry(max_retries=3, delay=1)
    async def index_document(self, index: str, doc: Dict[str, Any]) -> None:
        """
        单文档写入 Elasticsearch。
        """
        try:
            document = await self.client.index(index=index, document=doc)
            return document
        except Exception as e:
            logger.error(f"Failed to index document: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def bulk_index(self, index: str, docs: List[Dict[str, Any]]) -> None:
        """
        批量写入 Elasticsearch。
        优化版：正确构建批量操作请求格式
        """
        if not docs:
            return

        try:
            operations = []
            for doc in docs:
                operations.append({"index": {"_index": index}})
                operations.append(doc)
            await self.client.bulk(operations=operations)
        except Exception as e:
            logger.error(f"Failed to bulk index documents: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def get_document(self, index: str, doc_id: str) -> Dict[str, Any]:
        """
        根据文档 ID 获取单个文档。
        """
        try:
            result = await self.client.get(index=index, id=doc_id)
            return result['_source']
        except Exception as e:
            logger.error(f"Failed to get document with ID {doc_id}: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def get_document_has_id(self, index: str, doc_id: str):
        """
        根据文档 ID 获取单个文档。
        """
        try:
            result = await self.client.get(index=index, id=doc_id)

            return result['_id'], result['_source']
        except Exception as e:
            logger.error(f"Failed to get document with ID {doc_id}: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def search_documents(self, index: str, query: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        执行搜索查询，返回匹配的文档列表。
        """
        try:
            result = await self.client.search(index=index, body=query)
            hits = result['hits']['hits']
            return [hit['_source'] for hit in hits]
        except Exception as e:
            logger.error(f"Failed to search documents: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def search_documents_have_id(self, index: str, query: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        执行搜索查询，返回匹配的文档列表。
        """
        try:
            result = await self.client.search(index=index, body=query)
            hits = result['hits']['hits']
            return [{'_id': hit['_id'], 'source': hit['_source']} for hit in hits]
        except Exception as e:
            logger.error(f"Failed to search documents: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def search(self, index: str, query: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行搜索查询，返回匹配的文档列表。
        """
        try:
            result = await self.client.search(index=index, body=query)
            return result
        except Exception as e:
            logger.error(f"Failed to search documents: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def search_array_field(self, index: str, nested_path: str, query: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        查询嵌套数组字段（需字段类型为 nested）
        :param index: 索引名称
        :param nested_path: 嵌套字段路径（如 "city.districts" ）
        :param query: 嵌套层级的查询条件
        :return: 匹配的嵌套数组内容
        """
        try:
            es_query = {
                "query": {
                    "nested": {
                        "path": nested_path,
                        "query": query,
                        "inner_hits": {}  # 返回匹配的嵌套对象
                    }
                }
            }
            result = await self.client.search(index=index, body=es_query)
            return extract_nested_hits(result, nested_path)
        except Exception as e:
            logger.error(f"Failed to search nested array {nested_path}: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def update(self, index: str, _id: str, body: dict):
        try:
            await self.client.update(index=index, id=_id, body=body)
        except Exception as e:
            logger.error(f"Failed to update elastic document {index} with {body}: {e}")
            raise

    @retry(max_retries=3, delay=1)
    async def bulk(self, index: str, docs: List):
        success, failed = await async_bulk(self.client, gen_actions(docs, index), index=index)
        print(f"Successfully inserted {success} documents.")
        print(f"Failed to insert {failed} documents.")


def gen_actions(docs: list, index: str):
    for doc in docs:
        yield {
            "_op_type": "index",
            "_index": index,
            "_source": doc
        }


# 提取嵌套查询结果的工具函数
def extract_nested_hits(result: dict, path: str) -> list:
    """
    从嵌套查询结果中提取内部命中的文档
    :param result: Elasticsearch 查询结果
    :param path: 嵌套字段路径
    :return: 匹配的嵌套文档列表
    """
    hits = []
    for outer_hit in result['hits']['hits']:
        inner_hits = outer_hit.get('inner_hits', {}).get(path, {})
        hits.extend([hit['_source'] for hit in inner_hits.get('hits', [])])
    return hits


# 工厂函数，简化客户端获取
async def get_es_client() -> ESClient:
    """获取已初始化的ES客户端实例"""
    client = ESClient()
    await client.initialize()
    return client


T = TypeVar('T', bound=BaseModel)  # 泛型模型类型约束

from typing import Any, Dict, List, Optional, Type, TypeVar
from pydantic import BaseModel

T = TypeVar('T', bound=BaseModel)


class ESResponseWrapper:
    __slots__ = ('_raw', '_hits', '_aggs', '_meta', '_model_cls')  # 移除_docs，改为_hits

    def __init__(self, raw_response: dict, model_cls: Optional[Type[T]] = None):
        """
        :param raw_response: Elasticsearch原始响应
        :param model_cls: 可选文档模型类
        """
        self._raw = raw_response
        self._hits = None  # 存储完整的hit信息（含ID）
        self._aggs = None
        self._meta = None
        self._model_cls = model_cls

    # 核心优化：保留完整hit信息
    @property
    def hits(self) -> List[dict]:
        """返回完整命中数据（含ID和元数据）"""
        if self._hits is None:
            self._hits = self._raw.get("hits", {}).get("hits", [])
        return self._hits

    @property
    def docs(self) -> List[dict]:
        """向后兼容：只获取文档内容（_source）"""
        return [hit["_source"] for hit in self.hits]

    @property
    def doc_ids(self) -> List[str]:
        """专用ID获取方法（零额外内存开销）"""
        return [hit["_id"] for hit in self.hits]

        # 增强模型转换：保留ID信息

    def to_models(self) -> List[T]:
        """返回绑定ID的增强模型列表"""
        if not self._model_cls:
            raise ValueError("Model class not provided")

        return [
            self._create_model_with_id(hit)
            for hit in self.hits
        ]

    def _create_model_with_id(self, hit: dict) -> T:
        """创建携带ID的模型实例"""

        model = self._model_cls(**hit["_source"])
        # 动态注入ID属性（不修改原始模型定义）
        setattr(model, "es_id", hit["_id"])
        setattr(model, "es_index", hit["_index"])
        return model

    # 更新操作快捷方式
    def update_document(self, doc_id: str, update_body: dict,
                        es_client: Any, **kwargs) -> dict:
        """
        执行文档更新

        :param doc_id: 文档ID
        :param update_body: 更新内容（符合ES更新语法）
        :param es_client: Elasticsearch客户端实例
        :return: ES更新响应
        """
        # 查找对应的索引（从缓存中获取）
        index = next(
            (hit["_index"] for hit in self.hits if hit["_id"] == doc_id),
            None
        )
        if not index:
            raise ValueError(f"Document {doc_id} not found in response")

        return es_client.update(
            index=index,
            id=doc_id,
            body={"doc": update_body},
            **kwargs
        )

    # 元数据增强（包含ID统计）
    @property
    def meta(self) -> dict:
        """增强元数据（包含ID相关信息）"""
        if self._meta is None:
            self._meta = {
                "total": self._raw["hits"]["total"]["value"],
                "ids": self.doc_ids,  # 包含所有ID
                "first_id": self.hits[0]["_id"] if self.hits else None,
                "took_ms": self._raw["took"],
                "shards": self._raw["_shards"]
            }
        return self._meta

    # 聚合属性保持不变
    @property
    def aggs(self) -> dict:
        if self._aggs is None:
            self._aggs = self._raw.get("aggregations", {})
        return self._aggs